package com.github.zzay.common.core.utils;

import com.github.zzay.common.core.annotation.MyRateLimit;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Map;

/**
 * @author zzay
 * @className RateLimitAop
 * @description Rate limiting
 * @create 2022/06/14 17:18
 */
@Slf4j
@Aspect
@Component
public class RateLimitAop {

    /**
     * 不同接口对应不同的流量控制策略
     */
    private final Map<String, RateLimiter> limiterMap = Maps.newConcurrentMap();

    /**
     * 在 {@link MyRateLimit} 构造 around 类型切面处理
     *
     * @param joinPoint Join point
     */
    @Around("@annotation(com.github.zzay.common.core.annotation.MyRateLimit)")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        // Get the signature of the join point
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        // Get the pointed method
        Method method = signature.getMethod();
        // Get the customized annotation
        MyRateLimit rateLimit = method.getAnnotation(MyRateLimit.class);
        if (rateLimit != null) {
            String key = rateLimit.key();
            RateLimiter rateLimiter = null;
            if (!limiterMap.containsKey(key)) {
                // Create a rate limiter
                rateLimiter = RateLimiter.create(rateLimit.permitsPerSecond());
                limiterMap.put(key, rateLimiter);
                log.info("新建令牌桶 = {}，容量 = {}", key, rateLimit.permitsPerSecond());
            }
            rateLimiter = limiterMap.get(key);
            boolean acquire = rateLimiter.tryAcquire(rateLimit.timeout(), rateLimit.timeunit());
            if (!acquire) {
                log.debug("令牌桶 = {}，获取令牌失败", key);
                return null;
            }
        }
        return joinPoint.proceed();
    }

}
